# coding=utf-8
# 模型模块自定义
import numpy as np
import mindspore
from mindspore import Tensor
from AlexNet import AlexNet

x = Tensor(np.random.randn(1, 3, 224, 224), mindspore.float32)
network = AlexNet()
logits = network(x)
print(logits.shape)

# Parameter
# 参数(Parameter)是一类特殊的Tensor，是指在模型训练过程中可以对其值进行更新的变量。MindSpore提供mindspore.Parameter类进行Parameter的构造。为了对不同用途的Parameter进行区分，下面对两种不同类别的Parameter进行定义：
# 可训练参数。在模型训练过程中根据反向传播算法求得梯度后进行更新的Tensor，此时需要将required_grad设置为True。
# 不可训练参数。不参与反向传播，但需要更新值的Tensor（如BatchNorm中的mean和var变量），此时需要将requires_grad设置为False。

import numpy as np
import mindspore
from mindspore import nn
from mindspore import ops
from mindspore import Tensor, Parameter


class Network(nn.Cell):
    def __init__(self):
        super().__init__()
        self.w = Parameter(Tensor(np.random.randn(5, 3), mindspore.float32), name='w')  # weight
        self.b = Parameter(Tensor(np.random.randn(3, ), mindspore.float32), name='b')  # bias

    def construct(self, x):
        z = ops.matmul(x, self.w) + self.b
        return z


net = Network()
print(net.b.asnumpy())
# 获取可训练参数
print(net.trainable_params())
# 获取所有参数
print(type(net.get_parameters()))
for name, param in net.parameters_and_names():
    print(f"{name}:\n{param.asnumpy()}")

# 修改Parameter Parameter是一种特殊的Tensor，因此可以使用Tensor索引修改的方式对其值进行修改。
net.b[0] = 1.
print(net.b.asnumpy())
# 可调用Parameter.set_data方法，使用相同Shape的Tensor对Parameter进行覆盖。该方法常用于使用Initializer进行Cell遍历初始化。
net.b.set_data(Tensor([3, 4, 5]))
print(net.b.asnumpy())

# 运行时修改参数值
# 由于MindSpore的使用静态图加速编译设计，此时需要使用mindspore.ops.assign接口对参数进行赋值。该方法常用于自定义优化器场景。
import mindspore as ms


@ms.jit
def modify_parameter():
    b_hat = ms.Tensor([7, 8, 9])
    ops.assign(net.b, b_hat)
    return True


modify_parameter()
print(net.b.asnumpy())

# Parameter Tuple：变量元组ParameterTuple，用于保存多个Parameter，继承于元组tuple，提供克隆功能。
from mindspore.common.initializer import initializer
from mindspore import ParameterTuple

# 创建
x = Parameter(default_input=ms.Tensor(np.arange(2 * 3).reshape((2, 3))), name="x")
y = Parameter(default_input=initializer('ones', [1, 2, 3], ms.float32), name='y')
z = Parameter(default_input=2.0, name='z')
params = ParameterTuple((x, y, z))

# 从params克隆并修改名称为"params_copy"
params_copy = params.clone("params_copy")

print(params)
print(params_copy)

# Cell训练状态转换
# 神经网络中的部分Tensor操作在训练和推理时的表现并不相同，如nn.Dropout在训练时进行随机丢弃，但在推理时则不丢弃，
# nn.BatchNorm在训练时需要更新mean和var两个变量，在推理时则固定其值不变。因此我们可以通过Cell.set_train接口来设置神经网络的状态。
net.set_train()
print(net.phase)
# set_train(False)时，神经网络状态为predict：
net.set_train(False)
print(net.phase)


# 自定义神经网络层
# 可以根据需要，通过MindSpore提供的function接口、Primitive算子自定义神经网络层，并可以使用Cell.bprop方法自定义反向

# 使用function接口构造神经网络层
class Threshold(nn.Cell):
    def __init__(self, threshold, value):
        super().__init__()
        self.threshold = threshold
        self.value = value

    def construct(self, inputs):
        cond = ops.gt(inputs, self.threshold)
        value = ops.fill(inputs.dtype, inputs.shape, self.value)
        return ops.select(cond, inputs, value)


m = Threshold(0.1, 20)
inputs = mindspore.Tensor([0.1, 0.2, 0.3], mindspore.float32)
y = m(inputs)
print(y)

# 自定义Cell反向 :手动控制其反向的计算，此时我们可以通过Cell.bprop接口对其反向进行定义
class Dropout2d(nn.Cell):
    def __init__(self, keep_prob):
        super().__init__()
        self.keep_prob = keep_prob
        self.dropout2d = ops.Dropout2D(keep_prob)

    def construct(self, x):
        return self.dropout2d(x)
    # bprop方法分别有三个入参：
    # x: 正向输入，当正向输入为多个时，需同样数量的入参。
    # out: 正向输出。
    # dout: 反向传播时，当前Cell执行之前的反向结果。

    # 一般我们需要根据正向输出和前层反向结果配合，根据反向求导公式计算反向结果，并将其返回。
    # Dropout2d的反向计算需要根据正向输出的mask矩阵对前层反向结果进行mask，然后根据keep_prob进行缩放。最终可得到正确的计算结果。
    def bprop(self, x, out, dout):
        _, mask = out
        dy, _ = dout
        if self.keep_prob != 0:
            dy = dy * (1 / self.keep_prob)
        dy = mask.astype(mindspore.float32) * dy
        return (dy.astype(x.dtype), )

dropout_2d = Dropout2d(0.8)
dropout_2d.bprop_debug = True


# Hook功能
# 调试深度学习网络是每一个深度学习领域的从业者需要面对且投入精力较大的工作。由于深度学习网络隐藏了中间层算子的输入、输出数据以及反向梯度，
# 只提供网络输入数据（特征量、权重）的梯度，导致无法准确地感知中间层算子的数据变化，从而降低了调试效率。
# 为了方便用户准确、快速地对深度学习网络进行调试，MindSpore在动态图模式下设计了Hook功能，使用Hook功能可以捕获中间层算子的输入、输出数据以及反向梯度。
# 目前，动态图模式下提供了四种形式的Hook功能，分别是：HookBackward算子和在Cell对象上进行注册的register_forward_pre_hook、register_forward_hook、register_backward_hook功能。

# HookBackward算子
import mindspore as ms
from mindspore import ops

ms.set_context(mode=ms.PYNATIVE_MODE)

def hook_fn(grad_out):
    """打印梯度"""
    print("hook_fn print grad_out:", grad_out)

hook = ops.HookBackward(hook_fn)
def hook_test(x, y):
    z = x * y
    z = hook(z)
    z = z * y
    return z

def net(x, y):
    return ms.grad(hook_test, grad_position=(0, 1))(x, y)

output = net(ms.Tensor(1, ms.float32), ms.Tensor(2, ms.float32))
print("output:", output)

# Cell对象的register_forward_pre_hook功能
# 如果在Hook函数中对Cell对象的原始输入数据进行计算操作后，再作为新的输入数据返回，这些新增的计算操作将会同时作用于梯度的反向传播。
def forward_pre_hook_fn(cell_id, inputs):
    print("forward inputs: ", inputs)
import numpy as np
import mindspore as ms
import mindspore.nn as nn

ms.set_context(mode=ms.PYNATIVE_MODE)

def forward_pre_hook_fn(cell_id, inputs):
    print("forward inputs: ", inputs)
    input_x = inputs[0]
    return input_x

class Net(nn.Cell):
    def __init__(self):
        super(Net, self).__init__()
        self.relu = nn.ReLU()
        self.handle = self.relu.register_forward_pre_hook(forward_pre_hook_fn)

    def construct(self, x, y):
        x = x + y
        x = self.relu(x)
        return x

net = Net()
grad_net = ms.grad(net, grad_position=(0, 1))

x = ms.Tensor(np.ones([1]).astype(np.float32))
y = ms.Tensor(np.ones([1]).astype(np.float32))

output = net(x, y)
print(output)
gradient = grad_net(x, y)
print(gradient)
net.handle.remove()
gradient = grad_net(x, y)
print(gradient)

# 用户如果在Hook函数中直接返回新创建的数据，而不是返回由原始输入数据经过计算后得到的数据，那么梯度的反向传播将会在该Cell对象上截止。
import numpy as np
import mindspore as ms
import mindspore.nn as nn

ms.set_context(mode=ms.PYNATIVE_MODE)

def forward_pre_hook_fn(cell_id, inputs):
    print("forward inputs: ", inputs)
    return ms.Tensor(np.ones([1]).astype(np.float32))

class Net(nn.Cell):
    def __init__(self):
        super(Net, self).__init__()
        self.relu = nn.ReLU()
        self.handle = self.relu.register_forward_pre_hook(forward_pre_hook_fn)

    def construct(self, x, y):
        x = x + y
        x = self.relu(x)
        return x

net = Net()
grad_net = ms.grad(net, grad_position=(0, 1))

x = ms.Tensor(np.ones([1]).astype(np.float32))
y = ms.Tensor(np.ones([1]).astype(np.float32))

gradient = grad_net(x, y)
print(gradient)

# Cell对象的register_forward_hook功能
def forward_hook_fn(cell_id, inputs, outputs):
    print("forward inputs: ", inputs)
    print("forward outputs: ", outputs)

import numpy as np
import mindspore as ms
import mindspore.nn as nn

ms.set_context(mode=ms.PYNATIVE_MODE)

def forward_hook_fn(cell_id, inputs, outputs):
    print("forward inputs: ", inputs)
    print("forward outputs: ", outputs)
    outputs = outputs + outputs
    return outputs

class Net(nn.Cell):
    def __init__(self):
        super(Net, self).__init__()
        self.relu = nn.ReLU()
        self.handle = self.relu.register_forward_hook(forward_hook_fn)

    def construct(self, x, y):
        x = x + y
        x = self.relu(x)
        return x

net = Net()
grad_net = ms.grad(net, grad_position=(0, 1))

x = ms.Tensor(np.ones([1]).astype(np.float32))
y = ms.Tensor(np.ones([1]).astype(np.float32))

gradient = grad_net(x, y)
print(gradient)
net.handle.remove()
gradient = grad_net(x, y)
print(gradient)

# Cell对象的register_backward_hook功能
def backward_hook_function(cell_id, grad_input, grad_output):
    print(grad_input)
    print(grad_output)
import numpy as np
import mindspore as ms
import mindspore.nn as nn

ms.set_context(mode=ms.PYNATIVE_MODE)

def backward_hook_function(cell_id, grad_input, grad_output):
    print(grad_input)
    print(grad_output)

class Net(nn.Cell):
    def __init__(self):
        super(Net, self).__init__()
        self.conv = nn.Conv2d(1, 2, kernel_size=2, stride=1, padding=0, weight_init="ones", pad_mode="valid")
        self.bn = nn.BatchNorm2d(2, momentum=0.99, eps=0.00001, gamma_init="ones")
        self.handle = self.bn.register_backward_hook(backward_hook_function)
        self.relu = nn.ReLU()

    def construct(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x

net = Net()
grad_net = ms.grad(net)
output = grad_net(ms.Tensor(np.ones([1, 1, 2, 2]).astype(np.float32)))
print(output)
net.handle.remove()
output = grad_net(ms.Tensor(np.ones([1, 1, 2, 2]).astype(np.float32)))
print("-------------\n", output)
# register_backward_hook 中注册的Hook函数仅捕获原始Cell对象的输入、输出梯度。
import numpy as np
import mindspore as ms
import mindspore.nn as nn

ms.set_context(mode=ms.PYNATIVE_MODE)

def forward_pre_hook_fn(cell_id, inputs):
    print("forward inputs: ", inputs)
    input_x = inputs[0]
    return input_x

def forward_hook_fn(cell_id, inputs, outputs):
    print("forward inputs: ", inputs)
    print("forward outputs: ", outputs)
    outputs = outputs + outputs
    return outputs

def backward_hook_fn(cell_id, grad_input, grad_output):
    print("grad input: ", grad_input)
    print("grad output: ", grad_output)

class Net(nn.Cell):
    def __init__(self):
        super(Net, self).__init__()
        self.relu = nn.ReLU()
        self.handle = self.relu.register_forward_pre_hook(forward_pre_hook_fn)
        self.handle2 = self.relu.register_forward_hook(forward_hook_fn)
        self.handle3 = self.relu.register_backward_hook(backward_hook_fn)

    def construct(self, x, y):
        x = x + y
        x = self.relu(x)
        return x

net = Net()
grad_net = ms.grad(net, grad_position=(0, 1))
gradient = grad_net(ms.Tensor(np.ones([1]).astype(np.float32)), ms.Tensor(np.ones([1]).astype(np.float32)))
print(gradient)
